#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   train.py
@Author :   Song
@Time   :   2022/2/28 21:48
@Contact:   songjian@westlake.edu.cn
@intro  :
'''
import math
import operator
import einops
import click
import numba
from numba import jit
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch.optim
import torch.utils.data
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils
from pathlib import Path
import predifine
import AlphaNovo_models
from AlphaNovo_dataloader import Test_Dataset
import mgf
from pathlib import Path
from numba import cuda
import time

try:
    profile
except:
    profile = lambda x: x


# @profile
def test_one_epoch(dataloader, model_len, model_mat):
    model_len.eval()
    model_mat.eval()
    aa_to_mass = {}
    aa_to_mass['x'] = 0.
    for k, v in predifine.residues.items():
        aa_to_mass[k] = v + 18.010565
    aa_to_mass['$'] = 0.
    mass_idx = np.array(list(aa_to_mass.values()))
    pred_len_v, pred_mass_v, pred_idx_v = [], [], []
    finish_num = 0
    for batch_idx, (batch_spectra, batch_pr_charge) in enumerate(dataloader):
        batch_spectra = batch_spectra.float().to(predifine.device)
        batch_pr_charge = batch_pr_charge.long().to(predifine.device)
        # forward
        with torch.no_grad():
            # pred len
            pred_len = model_len(batch_spectra, batch_pr_charge)
            pred_len = torch.softmax(pred_len, dim=1)
            pred_len = torch.argmax(pred_len, dim=1) + 1
            pred_len_v.append(pred_len.type(torch.int8))

            # pred mat
            pred = model_mat(batch_spectra, batch_pr_charge)
            pred = einops.rearrange(pred, 'B (M L) -> B M L', M=predifine.max_length + 1)
            pred = torch.softmax(pred, dim=2)
            _, pred = torch.topk(pred, predifine.top_n, dim=2)
            pred_idx_v.append(pred.type(torch.int8))

            mass = torch.tensor(np.tile(mass_idx, (len(pred), predifine.max_length + 1, 1)), device=predifine.device)
            mass = torch.gather(mass, 2, pred)
            pred_mass_v.append(mass.type(torch.float32))

            finish_num += pred.shape[0]
            print('\rfinish: {}/{}'.format(finish_num, len(dataloader.dataset)), end='')
    print('\r')

    pred_len_v = torch.cat(pred_len_v)
    pred_idx_v = torch.cat(pred_idx_v, dim=0)
    pred_mass_v = torch.cat(pred_mass_v, dim=0)

    return pred_len_v, pred_idx_v, pred_mass_v


def my_collate(items):
    batch_spectra, batch_pr_charge = zip(*items)

    batch_spectra = rnn_utils.pad_sequence(batch_spectra, batch_first=True)
    batch_pr_charge = torch.tensor(batch_pr_charge)

    return batch_spectra, batch_pr_charge


@cuda.jit
def gpu_calculate_mass(n, pep_len, pred_mass,
                       mat_id, mass_left,
                       mass_right, step, possible_total,
                       result):
    # 每个线程的id对应处理一个查询项
    thread_idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    if thread_idx >= n:
        return
    start = thread_idx * step
    end = start + step

    cols = cuda.local.array(12, dtype=numba.int8)
    mat = pred_mass[mat_id]

    for i in range(start, end):
        cols[0] = i % 5
        for row in range(1, pep_len):
            col = i // (5 ** row) % 5
            cols[row] = col

        mass_cal = 0.
        for row in range(pep_len):
            mass = mat[row, cols[row]]
            if mass == 0.: # 出现了结束符，该idx无效
                return
            mass_cal += mat[row, cols[row]]
        if mass_left < mass_cal < mass_right:
            result[i] += 1


@cuda.jit
def gpu_idx_to_aa(n, pred_len, mat_id_v, pred_idx, candidate_idx_v, result):
    # 每个线程的id对应处理一个查询项
    thread_idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    if thread_idx >= n:
        return

    mat_id = mat_id_v[thread_idx]
    mat = pred_idx[mat_id]
    candidate = candidate_idx_v[thread_idx]

    for row in range(pred_len):
        if row == 0:
            col = candidate % 5
        else:
            col = candidate // (5 ** row) % 5
        aa_idx = mat[row, col]

        ''' 将aa_idx转换为对应字符的ascii码
        {1: 'G'-71, 2: 'A'-65, 3: 'S'-83, 4: 'P'-80, 5: 'V'-86, 
         6: 'T'-84, 7: 'c'-99, 8: 'L'-76, 9: 'I'-73,
         10: 'N'-78, 11: 'D'-68, 12: 'Q'-81, 13: 'K'-75, 14: 'E'-69, 15: 'M'-77, 
         16: 'H'-72, 17: 'F'-70, 18: 'R'-82, 19: 'Y'-89, 20: 'W'-87, 21: 'm'-109, 22: '$'-36}
        '''
        if aa_idx == 1:
            aa_ascii = 71
        elif aa_idx == 2:
            aa_ascii = 65
        elif aa_idx == 3:
            aa_ascii = 83
        elif aa_idx == 4:
            aa_ascii = 80
        elif aa_idx == 5:
            aa_ascii = 86
        if aa_idx == 6:
            aa_ascii = 84
        elif aa_idx == 7:
            aa_ascii = 99
        elif aa_idx == 8:
            aa_ascii = 76
        elif aa_idx == 9:
            aa_ascii = 73
        elif aa_idx == 10:
            aa_ascii = 78
        if aa_idx == 11:
            aa_ascii = 68
        elif aa_idx == 12:
            aa_ascii = 81
        elif aa_idx == 13:
            aa_ascii = 75
        elif aa_idx == 14:
            aa_ascii = 69
        elif aa_idx == 15:
            aa_ascii = 77
        if aa_idx == 16:
            aa_ascii = 72
        elif aa_idx == 17:
            aa_ascii = 70
        elif aa_idx == 18:
            aa_ascii = 82
        elif aa_idx == 19:
            aa_ascii = 89
        elif aa_idx == 20:
            aa_ascii = 87
        elif aa_idx == 21:
            aa_ascii = 109
        elif aa_idx == 22:
            aa_ascii = 36
        result[thread_idx, row] = aa_ascii


def cpu_calculate_mass(n, pep_len, pred_mass,
                       scan_id_v,
                       top_n, mass_left_v, mass_right_v,
                       threads_cumsum,
                       result):
    # 每个线程的id对应处理一个查询项
    thread_idx = 123456
    if thread_idx >= n:
        return

    thread_cumsum_idx = thread_idx + threads_cumsum

    for y, scan_id in enumerate(scan_id_v):
        mass_cal = 0.
        mat = pred_mass[scan_id]
        mass_left = mass_left_v[scan_id]
        mass_right = mass_right_v[scan_id]
        for row in range(pep_len):
            if row == 0:
                col = thread_cumsum_idx % 10
            else:
                col = thread_cumsum_idx // (10 ** row) % 10
            mass_cal += mat[row, col]
            if mass_cal > mass_right:
                break
        if mass_left < mass_cal < mass_right:
            result[y, thread_idx] = 1

    return result


@jit(nopython=True)
def idx_to_cols(idx, pep_len, top_n):
    cols = np.zeros(pep_len, dtype=np.int8)
    cols[0] = idx % top_n
    for row in range(1, pep_len):
        col = idx // (5 ** row) % 5
        cols[row] = col
    return cols


# @click.command()
# @click.option('--dir_mgf', required=False, help='Specify the target mgf folder dir')
def main():
    # train and valid
    predifine.device = torch.device('cuda:0')

    dir_mgf = r'/home/songjian/casanovo/test_mgf/GG20170112_CRH_HLA_C1203_biorep2_techrep1.mgf'
    dir_out = Path(r'/home/songjian/casanovo/test_mgf')
    dir_model_len = r'/home/songjian/casanovo/train_output_pred_len/alphanovo_epoch_9.pt'
    dir_model_mat = r'/home/songjian/casanovo/train_output_pred_mat/alphanovo_epoch_9.pt'

    print(dir_mgf, dir_model_len, dir_model_mat)

    ms = mgf.mgf_Reader(dir_mgf, is_annotated=False)
    ms.pr_charge[ms.pr_charge > predifine.max_charge] = predifine.max_charge

    df = pd.DataFrame()
    df['scan_id'] = np.arange(0, len(ms.pr_charge), dtype=np.uint32)
    df['scan_No'] = ms.scan_No
    df['pr_mz_measure'] = ms.pr_mz.astype(np.float32)
    df['pr_charge'] = ms.pr_charge.astype(np.int8)
    dataset = Test_Dataset('de novo', ms)
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=predifine.test_batch_size,
                                             num_workers=predifine.test_num_workers,
                                             shuffle=False,
                                             pin_memory=True,
                                             collate_fn=my_collate)  # 默认collate_fn

    # model
    model_len = AlphaNovo_models.AlphaNovo_Model_len(
        dim_model=predifine.dim_model,
        n_head=predifine.n_head,
        dim_feedforward=predifine.dim_feedforward,
        n_layers=predifine.n_layers,
        dropout=predifine.dropout,
        dim_intensity=predifine.dim_intensity,
        max_length=predifine.max_length,
        max_charge=predifine.max_charge)
    model_len.load_state_dict(torch.load(dir_model_len, map_location=predifine.device))
    model_len = model_len.to(predifine.device)
    model_mat = AlphaNovo_models.AlphaNovo_Model_mat(
        dim_model=predifine.dim_model,
        n_head=predifine.n_head,
        dim_feedforward=predifine.dim_feedforward,
        n_layers=predifine.n_layers,
        dropout=predifine.dropout,
        dim_intensity=predifine.dim_intensity,
        max_length=predifine.max_length,
        max_charge=predifine.max_charge)
    model_mat.load_state_dict(torch.load(dir_model_mat, map_location=predifine.device))
    model_mat = model_mat.to(predifine.device)

    t0 = time.time()
    pred_len, pred_idx, pred_mass = test_one_epoch(dataloader, model_len, model_mat)
    assert len(pred_mass) == len(pred_idx) == len(df)
    print(f'模型预测耗时：{(time.time() - t0):.3f}s')

    df['pred_len'] = pred_len.cpu().numpy()
    df = df[(df['pred_len'] > 6) & (df['pred_len'] < 13)]
    # df = df.sort_values(by='pred_len', ascending=True)
    # df = df.sample(1000)
    df = df.reset_index(drop=True)
    mass_H, mass_h2o = 1.007276466771, 18.0105650638
    aa_mass_sum = (df['pr_mz_measure'] * df['pr_charge'] - mass_H * df['pr_charge'] \
                   + mass_h2o * (df['pred_len'] - 1)).astype(np.float32)
    mass_left_v = (aa_mass_sum * (1 - predifine.ppm / 1000000.)).astype(np.float32)
    mass_right_v = (aa_mass_sum * (1 + predifine.ppm / 1000000.)).astype(np.float32)
    pred_len_v = df['pred_len'].to_numpy()

    # 肽段分批次、枚举分批次计算每个肽段的有效候选
    pred_mass = cuda.as_cuda_array(pred_mass)
    pred_idx = cuda.as_cuda_array(pred_idx)
    scan_id_v = df['scan_id'].to_numpy()

    result_idx_v = []
    t0 = time.time()
    for i in range(len(df)):
        if i % 10000 == 0:
            print(i)
        pred_len = pred_len_v[i]
        mat_id = scan_id_v[i]
        mass_left = mass_left_v[i]
        mass_right = mass_right_v[i]
        possible_total = predifine.top_n ** pred_len
        n = 5 ** 9
        step = math.ceil(possible_total / n)  # 每一个thread负责这么多候选
        if step == 1:
            n = possible_total
        threads_per_block = 512
        blocks_per_grid = math.ceil(n / threads_per_block)
        result = cuda.to_device(np.zeros(possible_total, dtype=np.uint8))
        # t0 = time.time()
        gpu_calculate_mass[blocks_per_grid, threads_per_block](n, pred_len, pred_mass,
                                                               mat_id, mass_left,
                                                               mass_right, step, possible_total,
                                                               result)
        cuda.synchronize()
        result = result.copy_to_host()
        result = np.where(result > 0)[0]
        result_idx_v.append(result)

    df['candidate_idx'] = result_idx_v
    df['candidate_num'] = df['candidate_idx'].apply(len)
    df = df[df['candidate_num'] > 0]
    candidate_num = df['candidate_num'].to_numpy()
    print(f'计算候选符合ppm：{(time.time() - t0):.3f}s')

    df_result = pd.DataFrame()
    df_result['scan_id'] = df['scan_id'].repeat(candidate_num)
    df_result['pr_mz_measure'] = df['pr_mz_measure'].repeat(candidate_num)
    df_result['pr_charge'] = df['pr_charge'].repeat(candidate_num)
    df_result['pred_len'] = df['pred_len'].repeat(candidate_num)

    df_result['candidate_idx'] = df['candidate_idx'].explode().to_numpy().astype(np.uint32)

    df_v, result_v = [], []
    t0 = time.time()
    for pred_len in df_result['pred_len'].unique():
        print(pred_len)
        df_len = df_result[df_result['pred_len'] == pred_len]
        mat_id_v = cuda.to_device(df_len['scan_id'].to_numpy())
        candidate_idx_v = cuda.to_device(df_len['candidate_idx'].to_numpy())

        n = len(df_len)
        threads_per_block = 512
        blocks_per_grid = math.ceil(n / threads_per_block)
        result = cuda.to_device(np.zeros((n, pred_len), dtype=np.uint8))
        gpu_idx_to_aa[blocks_per_grid, threads_per_block](n, pred_len, mat_id_v,
                                                          pred_idx, candidate_idx_v, result)
        cuda.synchronize()
        result = result.copy_to_host()
        result = result.view('S' + str(pred_len)).squeeze()
        df_len.loc[:, 'seq'] = result.astype('U' + str(pred_len))
        df_v.append(df_len)
    print(f'结果转换为seq：{(time.time() - t0):.3f}s')

    t0 = time.time()
    df_result = pd.concat(df_v)
    del df_result['candidate_idx']

    # df_result.to_pickle(dir_out / 'gg.pkl')
    # df_result.to_hdf(dir_out / 'gg.h5', 'data')
    df_result.to_csv(dir_out / 'gg.tsv', sep='\t', index=False)

    print(f'{len(df_result)}保存用时：{(time.time() - t0):.3f}s')

if __name__ == '__main__':
    t0 = time.time()
    main()
    print(f'Finished. time: {(time.time() - t0):.3f}')
